package dicebot.apps;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.URL;
import java.net.URLConnection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.UUID;

import dicebot.abstracts.MessageHandler;
import dicebot.filene.Main;
import net.mamoe.mirai.event.events.GroupMessageEvent;
import net.mamoe.mirai.message.data.At;
import net.mamoe.mirai.message.data.MessageChain;
import net.mamoe.mirai.message.data.MessageChainBuilder;
import net.sf.json.JSONArray;
import net.sf.json.JSONObject;

public class ChatGPT {
	
	public static final Map<String, String> headers = new HashMap<>();
	public static void readHeaders() throws Exception {
		File file = new File("header.txt");
		BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(file)));
		String line;
		while ((line = reader.readLine()) != null) {
			if (line.startsWith(":"))
				continue;
			int i = line.indexOf(":", 1);
			String key = line.substring(0, i);
			String value = line.substring(i + 1, line.length());
			headers.put(key, value);
		}
		reader.close();
	}

	public static void initialize () throws Exception {
		readHeaders();
		
		Statement statement = Sql.getStatement();
		statement.execute("create table if not exists ChatGPT (QQ text not null primary key, ConversationId text, LastMessageId text);");
		statement.close();
		
		Main.HANDLERS.add(new MessageHandler() {
			@Override
			public boolean handle(final GroupMessageEvent e) throws Exception {
				String msg = e.getMessage().contentToString();
				if (!(msg.startsWith(":") || msg.startsWith("："))) {
					return false;
				}
				final String colon = msg.substring(0, 1);
				msg = msg.substring(1);
				final long sender = e.getSender().getId();
				Statement statement = Sql.getStatement();
				ResultSet r = statement.executeQuery("select * from ChatGPT where QQ = '" + sender + "';");
				boolean record = r.next();
				
				URL url = new URL("https://chat.openai.com/backend-api/conversation");
				URLConnection connection = url.openConnection();

				// load headers
				for (Entry<String, String> entry : headers.entrySet()) {
					connection.setRequestProperty(entry.getKey(), entry.getValue());
				}
				
				connection.setDoOutput(true);
				connection.connect();
				OutputStream out = connection.getOutputStream();

				// send pay load
				JSONObject obj = new JSONObject();
				obj.put("action", "next");
				obj.put("model", "text-davinci-002-render");
				if (record && r.getString("LastMessageId") != null) {
					obj.put("parent_message_id", r.getString("LastMessageId"));
				} else {
					obj.put("parent_message_id", UUID.randomUUID().toString());
				}
				if (record && r.getString("ConversationId") != null) {					
					obj.put("conversation_id", r.getString("ConversationId"));
				}
				JSONArray messages = new JSONArray();
				JSONObject message = new JSONObject();
				message.put("id", UUID.randomUUID().toString());
				message.put("role", "user");
				JSONObject msgContent = new JSONObject();
				msgContent.put("content_type", "text");
				JSONArray msgParts = new JSONArray();
				msgParts.add(msg);
				msgContent.put("parts", msgParts);
				message.put("content", msgContent);
				messages.add(message);
				obj.put("messages", messages);
				statement.close();
				System.out.println(obj.toString());
				out.write(obj.toString().getBytes());
				out.flush();
				
				// wait for get response
				InputStream in = connection.getInputStream();
				final BufferedReader reader = new BufferedReader(new InputStreamReader(in));
				new Thread() {
					public void run() {
						try {							
							String line = null;
							String answer = null;
							while((line = reader.readLine()) != null) {
								if (line.isEmpty()) continue;
								if (line.equals("data: [DONE]")) break;
								answer = line;
							}
							JSONObject object = JSONObject.fromObject("{" + answer + "}").getJSONObject("data");
							String conversationId = object.getString("conversation_id");
							String messageId = object.getJSONObject("message").getString("id");
							JSONArray texts = object.getJSONObject("message").getJSONObject("content").getJSONArray("parts");
							String text = "";
							for (Object obj : texts) {
								text += obj;
							}
							MessageChain chain = new MessageChainBuilder()
									.append(new At(e.getSender().getId()))
									.append(colon + text).build();
							synchronized (ChatGPT.class) {								
								e.getGroup().sendMessage(chain);
							
								// store next message
								Statement statement = Sql.getStatement();
								statement.execute("insert or replace into ChatGPT (QQ, ConversationId, LastMessageId) values ('" + sender + "', '" + conversationId + "', '" + messageId + "');");
								statement.close();
							}
							
						} catch (Exception e2) {
							e2.printStackTrace();
						}
					};
				}.start();
				return true;
			}
		});
		Main.HANDLERS.add(new MessageHandler() {
			@Override
			public boolean handle(GroupMessageEvent e) throws Exception {
				String msg = e.getMessage().contentToString();
				if (!msg.startsWith(".reset")) {
					return false;
				}
				final long sender = e.getSender().getId();
				Statement statement = Sql.getStatement();
				ResultSet r = statement.executeQuery("select * from ChatGPT where QQ = '" + sender + "';");
				String conversationId = r.getString("ConversationId");
				String lastMessageId = r.getString("LastMessageId");
				if (r.next()) {
					statement.execute("delete from ChatGPT where QQ = '" + sender + "';");
					MessageChain chain = new MessageChainBuilder()
							.append(new At(e.getSender().getId()))
							.append("chatGPT线程已重置。如果需要恢复刚刚被重置的线程，请使用conversationId=" + conversationId + ", lastMessageId=" + lastMessageId + "来恢复。").build();
					e.getGroup().sendMessage(chain);
				} else {
					MessageChain chain = new MessageChainBuilder()
							.append(new At(e.getSender().getId()))
							.append("没有需要重置的chatGPT线程。").build();
					e.getGroup().sendMessage(chain);
				}
				statement.close();
				return false;
			}
		});
	}
	
}
